import sqlite3
from pathlib import Path
import sys
import argparse
from datetime import datetime
import zipfile
import os  # 导入os模块

def check_and_connect_db(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 设置返回字典格式的结果
    cursor = conn.cursor()
    return conn, cursor

def process_zip_file(zip_path, db_directory, file_path, file_name, file_hash, cursor, conn):
    zip_name = zip_path.name
    zip_rel_path = str(zip_path.relative_to(db_directory))
    
    create_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    
    # 插入到photo_hash表中
    cursor.execute("""
        INSERT INTO photo_hash (zip_path, zip_name, file_path, file_name, file_hash, create_time)
        VALUES (?, ?, ?, ?, ?, ?)
    """, (zip_rel_path, zip_name, file_path, file_name, file_hash, create_time))
    print(f"已插入photo_hash记录: {file_name} 在 {zip_name}，哈希值: {file_hash}")
    
    conn.commit()

def process_records(db_directory, db_conn, cursor):
    print("开始处理 deal_flag=3 的记录")
    
    # 查询photo_move表中所有deal_flag为3的记录
    cursor.execute("SELECT * FROM photo_move WHERE deal_flag=3 ORDER BY id ASC")  # 假设id字段用于排序
    records = cursor.fetchall()

    if not records:
        print("没有需要处理的记录")
        return
    
    for record in records:
        aim_zip = record['aim_zip']
        zip_path = Path(db_directory) / aim_zip
        
        if zip_path.exists() and zipfile.is_zipfile(zip_path):
            file_hash = record['file_hash']  # 直接从photo_move表中获取file_hash
            file_path = record['file_path']
            file_name = record['file_name']
            process_zip_file(zip_path, db_directory, file_path, file_name, file_hash, cursor, db_conn)
        else:
            print(f"警告: ZIP文件 {zip_path} 不存在或不是有效的ZIP文件")
            continue
        
        # 更新此记录的deal_flag为4
        cursor.execute("UPDATE photo_move SET deal_flag=4 WHERE id=?", (record['id'],))
        db_conn.commit()
        print(f"已更新deal_flag为4: 记录ID {record['id']}")

def main(db_directory):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    process_records(db_directory, conn, cursor)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process records with deal_flag=3, write to photo_hash and update deal_flag.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hashes.db)")

    args = parser.parse_args()

    db_directory = args.db

    if not os.path.isdir(db_directory):
        print(f"错误: {db_directory} 不是一个有效的目录")
        sys.exit(1)

    main(db_directory)